from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import openai
import numpy as np
import time

# ==== 配置 ====
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"
COLLECTION_NAME = "semantic_search"
EMBEDDING_DIM = 1536  # 取决于使用的模型，如 text-embedding-3-small 是 1536
openai.api_key = "your-openai-api-key"

# ==== 连接 Milvus ====
connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)

# ==== 创建 Collection ====
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=1000),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=EMBEDDING_DIM)
]
schema = CollectionSchema(fields, description="Semantic search demo")

if COLLECTION_NAME not in [c.name for c in Collection.list()]:
    collection = Collection(name=COLLECTION_NAME, schema=schema)
else:
    collection = Collection(name=COLLECTION_NAME)

# ==== 插入文本 + 向量 ====
def get_embedding(text):
    response = openai.Embedding.create(
        model="text-embedding-3-small",
        input=text
    )
    return response["data"][0]["embedding"]

texts = [
    "如何管理现场外包人员？",
    "什么是项目交付保障机制？",
    "如何做好软件测试？",
    "云计算的基本概念是什么？"
]

ids = [int(time.time()) + i for i in range(len(texts))]
embeddings = [get_embedding(text) for text in texts]

collection.insert([ids, texts, embeddings])
collection.load()

# ==== 搜索：语义相似问题 ====
query = "如何做好人力资源管理？"
query_embedding = get_embedding(query)

results = collection.search(
    data=[query_embedding],
    anns_field="embedding",
    param={"metric_type": "COSINE", "params": {"nprobe": 10}},
    limit=3,
    output_fields=["content"]
)

# ==== 输出结果 ====
print(f"\n🔍 查询: {query}")
for hit in results[0]:
    print(f"➡️ 相似内容: {hit.entity.get('content')} (相似度: {hit.distance:.4f})")
